
#include "opencv2/opencv.hpp"
#import "pixelBufferPool.h"
    
PixelBufferPool::PixelBufferPool() : pool_(NULL){
}
PixelBufferPool::~PixelBufferPool(){
    ReleaseCVPixelBufferPool(pool_);
}
    
std::shared_ptr<cv::Mat> PixelBufferPool::GetMat(CVPixelBufferRef pixelBuffer){
    size_t width = CVPixelBufferGetWidth(pixelBuffer);
    size_t height = CVPixelBufferGetHeight(pixelBuffer);
    auto mat = std::make_shared<cv::Mat>((int)height, (int)width, CV_8UC4);
    size_t channel = mat->channels();
    
    CVPixelBufferLockBaseAddress(pixelBuffer, 0);
    void *address = CVPixelBufferGetBaseAddress(pixelBuffer);
    if (width * channel * sizeof(uchar) == CVPixelBufferGetBytesPerRow(pixelBuffer)) {
        memcpy((uchar*)mat->data, (uchar*)address, width * height * channel *sizeof(uchar));
    } else {
        size_t perRow = CVPixelBufferGetBytesPerRow(pixelBuffer);
        for (int j = 0; j < height; ++j){
            //size_t perRow = CVPixelBufferGetBytesPerRowOfPlane(pixelBuffer, j);
            memcpy((uchar*)mat->data + perRow * j/sizeof(uchar), (uchar*)address + width*channel * j, perRow);
        }
    }
    CVPixelBufferUnlockBaseAddress(pixelBuffer, 0);
//    cv::cvtColor(*mat, *mat, CV_BGRA2BGR);
    return mat;
}
    
CVPixelBufferRef PixelBufferPool::GetPixelBuffer(const cv::Mat &mat){
    size_t channel = mat.channels();
    size_t width = mat.cols;
    size_t height = mat.rows;
//    cv::cvtColor(mat, mat, CV_BGR2BGRA);
    if (!pool_){
        pool_ = CreateCVPixelBufferPool((int)width, (int)height);
    }
    CVPixelBufferRef pixelBuffer = NULL;
    CVPixelBufferPoolCreatePixelBuffer(NULL, pool_, &pixelBuffer);
    if (pixelBuffer){
        CVPixelBufferLockBaseAddress(pixelBuffer, 0);
        void *address = CVPixelBufferGetBaseAddress(pixelBuffer);
        if (width * channel * sizeof(uchar) == CVPixelBufferGetBytesPerRow(pixelBuffer)) {
            memcpy((uchar*)address, (uchar*)mat.data, width * height * channel* sizeof(uchar));
        } else {
            size_t perRow = CVPixelBufferGetBytesPerRow(pixelBuffer);
            for (int j = 0; j < height; ++j){
                memcpy((uchar*)address + perRow * j/sizeof(uchar), (uchar*)mat.data + width * channel * j, perRow);
            }
        }
        CVPixelBufferUnlockBaseAddress(pixelBuffer, 0);
    }
    return pixelBuffer;
}
 
    CVPixelBufferPoolRef PixelBufferPool:: CreateCVPixelBufferPool(int width, int height){
        CVPixelBufferPoolRef pool = NULL;
        NSMutableDictionary *attributes = [NSMutableDictionary dictionary];
        [attributes setObject:[NSNumber numberWithInt:kCVPixelFormatType_32BGRA] forKey:(NSString*)kCVPixelBufferPixelFormatTypeKey];
        [attributes setObject:[NSNumber numberWithInt:width] forKey: (NSString*)kCVPixelBufferWidthKey];
        [attributes setObject:[NSNumber numberWithInt:height] forKey: (NSString*)kCVPixelBufferHeightKey];
        CVReturn theError = CVPixelBufferPoolCreate(kCFAllocatorDefault, NULL, (__bridge CFDictionaryRef)attributes, &pool);
        if ( !(kCVReturnSuccess == theError) ) {
            NSLog(@"offscreen render create pixel pool failed!");
        }
        return pool;
    }

    

void PixelBufferPool::ReleaseCVPixelBufferPool(CVPixelBufferPoolRef pool){
    if (pool) {
        CVPixelBufferPoolRelease(pool);
    }
}


void PixelBufferPool::SaveMatToFile(const cv::Mat &mat, NSString *filePath){
    NSArray *paths = NSSearchPathForDirectoriesInDomains(NSDocumentDirectory, NSUserDomainMask, YES);
    NSString *file = [paths[0] stringByAppendingPathComponent:filePath];
    if(!cv::imwrite(file.UTF8String, mat)){
        NSLog(@"write mat file failed!");
    }
}
